// Created by Cal Stephens on 9/25/24.
// Copyright © 2024 Airbnb Inc. All rights reserved.

public extension FormatRule {
    static let redundantEquatable = FormatRule(
        help: "Omit a hand-written Equatable implementation when the compiler-synthesized conformance would be equivalent.",
        disabledByDefault: true,
        options: ["equatablemacro"]
    ) { formatter in
        // Find all of the types with an `Equatable` conformance and a manually-implemented `static func ==` implementation.
        let declarations = formatter.parseDeclarations()
        let typesManuallyImplementingEquatableConformance = formatter.manuallyImplementedEquatableTypes(in: declarations)

        for equatableType in typesManuallyImplementingEquatableConformance {
            let isEligibleForAutoEquatableConformance: Bool
            switch equatableType.typeDeclaration.keyword {
            case "struct":
                // The compiler automatically synthesizes Equatable implementations for structs
                isEligibleForAutoEquatableConformance = true
            case "class":
                // Projects can define an `@Equatable` macro that generates the Equatable implementation for classes
                isEligibleForAutoEquatableConformance = formatter.options.equatableMacro != .none
            default:
                // This rule doesn't support other kinds of types.
                isEligibleForAutoEquatableConformance = false
            }

            guard isEligibleForAutoEquatableConformance,
                  let typeBody = equatableType.typeDeclaration.body
            else { continue }

            // Find all of the stored instance properties in this type.
            // The synthesized Equatable implementation would compare each of these.
            let storedInstanceProperties = typeBody.filter(\.isStoredInstanceProperty)
            let storedInstancePropertyNames = Set(storedInstanceProperties.map(\.name))

            // Find all of the properties compared using `lhs.{property} == rhs.{property}`
            let comparedProperties = formatter.parseComparedProperties(inEquatableImplementation: equatableType.equatableFunction)

            // If the set of compared properties match the set of stored instance properties,
            // then the manually implemented `==` function is redundant and can be removed.
            guard comparedProperties == storedInstancePropertyNames else {
                continue
            }

            // The compiler automatically synthesizes Equatable implementations for structs
            // as long as all of the properties are themselves Equatable. This is usually true
            //
            if equatableType.typeDeclaration.keyword == "struct",
               !storedInstanceProperties.contains(where: { $0.parsePropertyDeclaration()?.type?.name.isKnownNonEquatableType == true })
            {
                equatableType.equatableFunction.remove()
            }

            // In projects using an `@Equatable` macro, the Equatable implementation
            // can be generated by that macro instead of written manually.
            else if case let .macro(macro, module: module) = formatter.options.equatableMacro {
                let declarationWithEquatableConformance = equatableType.declarationWithEquatableConformance

                guard let equatableConformance = formatter.parseConformancesOfType(atKeywordIndex: declarationWithEquatableConformance.keywordIndex).first(where: { $0.conformance == "Equatable" || $0.conformance == "Hashable" })
                else { continue }

                // Exclude cases where the Equatable conformance is defined in an extension with a where clause,
                // since this wouldn't usually be captured in the generated conformance.
                if let startOfExtensionTypeBody = formatter.index(of: .startOfScope("{"), after: equatableConformance.index),
                   formatter.index(of: .keyword("where"), in: equatableConformance.index ..< startOfExtensionTypeBody) != nil
                {
                    continue
                }

                // Remove the `==` implementation
                equatableType.equatableFunction.remove()

                // Remove the `: Equatable` conformance.
                //  - If this type uses as `: Hashable` conformance, we have to preserve that.
                if equatableConformance.conformance == "Equatable" {
                    formatter.removeConformance(at: equatableConformance.index)
                }

                // Add the `@Equatable` macro
                formatter.insert(
                    [.keyword(macro), .space(" ")],
                    at: equatableType.typeDeclaration.startOfModifiersIndex
                )

                // Import the module that defines the `@Equatable` macro if needed
                formatter.addImports([module])
            }
        }
    } examples: {
        """
        ```diff
          struct Foo: Equatable {
              let bar: Bar
              let baaz: Baaz

        -     static func ==(lhs: Foo, rhs: Foo) -> Bool {
        -         lhs.bar == rhs.bar 
        -             && lhs.baaz == rhs.baaz
        -     }
          }

          class Bar: Equatable {
              let baaz: Baaz

              static func ==(lhs: Bar, rhs: Bar) -> Bool {
                  lhs.baaz == rhs.baaz
              }
          }
        ```

        If your project includes a macro that generates the `static func ==` implementation
        for the attached class, you can specify `--equatablemacro @Equatable,MyMacroLib`
        and this rule will also migrate eligible classes to use your macro instead of
        a hand-written Equatable conformance:

        ```diff
          // --equatablemacro @Equatable,MyMacroLib
          import FooLib
        + import MyMacroLib

        + @Equatable
        + class Bar {
        - class Bar: Equatable {
              let baaz: Baaz

        -     static func ==(lhs: Bar, rhs: Bar) -> Bool {
        -         lhs.baaz == rhs.baaz
        -     }
          }
        ```
        """
    }
}

extension Formatter {
    struct EquatableType {
        /// The main type declaration of the type that has an Equatable conformance
        let typeDeclaration: Declaration
        /// The Equatable `static func ==` implementation, which could be defined in an extension.
        let equatableFunction: Declaration
        /// The declaration that contains the `: Equatable` conformance, which may be an extension.
        let declarationWithEquatableConformance: Declaration
    }

    /// Finds all of the types in the current file with an Equatable conformance,
    /// which also have a manually-implemented `static func ==` method.
    func manuallyImplementedEquatableTypes(in declarations: [Declaration]) -> [EquatableType] {
        var typeDeclarationsByFullyQualifiedName: [String: Declaration] = [:]
        var typesWithEquatableConformances: [(fullyQualifiedTypeName: String, declarationWithEquatableConformance: Declaration)] = []
        var equatableImplementationsByFullyQualifiedName: [String: Declaration] = [:]

        declarations.forEachRecursiveDeclaration { declaration in
            guard let declarationName = declaration.name else { return }

            if declaration.definesType, let fullyQualifiedName = declaration.fullyQualifiedName {
                typeDeclarationsByFullyQualifiedName[fullyQualifiedName] = declaration
            }

            // Support the Equatable conformance being declared in an extension
            // separately from the Equatable
            if declaration is TypeDeclaration,
               let fullyQualifiedName = declaration.fullyQualifiedName
            {
                let conformances = parseConformancesOfType(atKeywordIndex: declaration.keywordIndex)

                // Both an Equatable and Hashable conformance will cause the Equatable conformance to be synthesized
                if conformances.contains(where: {
                    $0.conformance == "Equatable" || $0.conformance == "Hashable"
                }) {
                    typesWithEquatableConformances.append((
                        fullyQualifiedTypeName: fullyQualifiedName,
                        declarationWithEquatableConformance: declaration
                    ))
                }
            }

            if declaration.keyword == "func",
               declarationName == "==",
               modifiersForDeclaration(at: declaration.keywordIndex, contains: "static"),
               let startOfArguments = index(of: .startOfScope("("), after: declaration.keywordIndex)
            {
                let functionArguments = parseFunctionDeclarationArguments(startOfScope: startOfArguments)

                if functionArguments.count == 2,
                   // The external label doesn't matter, it can be `_` or `lhs/rhs`.
                   functionArguments[0].internalLabel == "lhs",
                   functionArguments[1].internalLabel == "rhs",
                   functionArguments[0].type == functionArguments[1].type
                {
                    var comparedTypeName = functionArguments[0].type

                    if let parentDeclaration = declaration.parent {
                        // If the function uses `Self`, resolve that to the name of the parent type
                        if comparedTypeName == "Self",
                           let parentDeclarationName = parentDeclaration.fullyQualifiedName
                        {
                            comparedTypeName = parentDeclarationName
                        }

                        // If the function uses `Bar` in an extension `Foo.Bar`, then resolve
                        // the name of the compared type to be the fully-qualified `Foo.Bar` type.
                        if parentDeclaration.keyword == "extension",
                           let extendedType = parentDeclaration.name,
                           comparedTypeName != extendedType,
                           extendedType.hasSuffix("." + comparedTypeName)
                        {
                            comparedTypeName = extendedType
                        }

                        // If the function uses `Bar` in a type `Bar`, then resolve the
                        // the name of the compared type to be the fully-qualified parent type.
                        //  - For example, `Bar` could be defined in a parent `Foo` type.
                        if comparedTypeName == parentDeclaration.name,
                           let parentDeclarationName = parentDeclaration.fullyQualifiedName
                        {
                            comparedTypeName = parentDeclarationName
                        }
                    }

                    equatableImplementationsByFullyQualifiedName[comparedTypeName] = declaration
                }
            }
        }

        return typesWithEquatableConformances.compactMap { typeName, declarationWithEquatableConformance in
            guard let typeDeclaration = typeDeclarationsByFullyQualifiedName[typeName],
                  let equatableImplementation = equatableImplementationsByFullyQualifiedName[typeName]
            else { return nil }

            return EquatableType(
                typeDeclaration: typeDeclaration,
                equatableFunction: equatableImplementation,
                declarationWithEquatableConformance: declarationWithEquatableConformance
            )
        }
    }

    /// Finds the set of properties that are compared in the given Equatable `func`,
    /// following the pattern `lhs.{property} == rhs.{property}`.
    ///  - Returns `nil` if there are any comparisons that don't match this pattern.
    func parseComparedProperties(inEquatableImplementation equatableImplementation: Declaration) -> Set<String>? {
        let funcIndex = equatableImplementation.keywordIndex

        guard let startOfBody = index(of: .startOfScope("{"), after: funcIndex),
              let firstIndexInBody = index(of: .nonSpaceOrCommentOrLinebreak, after: startOfBody),
              let endOfBody = endOfScope(at: startOfBody)
        else { return nil }

        var validComparedProperties = Set<String>()
        var currentIndex = firstIndexInBody

        // Skip over any `return` keyword that may be present
        if tokens[currentIndex] == .keyword("return"),
           let nextIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: currentIndex)
        {
            currentIndex = nextIndex
        }

        while currentIndex < endOfBody {
            // Parse the current `lhs.{property} == rhs.{property}` pattern
            guard tokens[currentIndex] == .identifier("lhs"),
                  let lhsDotIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: currentIndex),
                  tokens[lhsDotIndex] == .operator(".", .infix),
                  let lhsPropertyName = index(of: .nonSpaceOrCommentOrLinebreak, after: lhsDotIndex),
                  tokens[lhsPropertyName].isIdentifierOrKeyword,
                  let equalsIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: lhsPropertyName),
                  tokens[equalsIndex] == .operator("==", .infix),
                  let rhsIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: equalsIndex),
                  tokens[rhsIndex] == .identifier("rhs"),
                  let rhsDotIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: rhsIndex),
                  tokens[rhsDotIndex] == .operator(".", .infix),
                  let rhsPropertyName = index(of: .nonSpaceOrCommentOrLinebreak, after: rhsDotIndex),
                  tokens[rhsPropertyName] == tokens[lhsPropertyName],
                  let indexAfterComparison = index(of: .nonSpaceOrCommentOrLinebreak, after: rhsPropertyName)
            else {
                // If we find a non-matching comparison, we have to avoid modifying this declaration
                return nil
            }

            validComparedProperties.insert(tokens[lhsPropertyName].string)

            // Skip over any `&&` operators connecting two comparisons
            if tokens[indexAfterComparison] == .operator("&&", .infix),
               let indexAfterAndOperator = index(of: .nonSpaceOrCommentOrLinebreak, after: indexAfterComparison)
            {
                currentIndex = indexAfterAndOperator
            }

            else {
                currentIndex = indexAfterComparison
            }
        }

        return validComparedProperties
    }
}

extension String {
    /// Whether or not this type name is known to be non-Equatable
    var isKnownNonEquatableType: Bool {
        let knownNonEquatableTypes = ["AnyClass"]
        return knownNonEquatableTypes.contains(self) || isTupleType
    }
}
